from abc import ABC, abstractmethod
from stable_baselines3 import PPO
from policy_manger import PolicyManager
from typing import Dict,List
# 子策略基类 定义子策略的基本接口和结构
class SubStrategy(ABC):
    _registry = {}      # 子策略类型注册表

    def __init_subclass__(cls,strategy_type=None, **kwargs):
        super().__init_subclass__(**kwargs)
        if strategy_type is not None:
            cls._registry[strategy_type] = cls

    def __init__(self,strategy_type):
        self.strategy_type = strategy_type
        self.policy_manager = PolicyManager(sub_strategy_types=strategy_type)               # 管理子策略的策略模型
        self.current_policy = self.policy_manager.get_latest_policy()                       # 当前策略
        
    @classmethod
    def creat_strategy(cls, strategy_type,aircraft_id=None):
        strategy_class = cls._registry.get(strategy_type)
        if strategy_class is None:
            raise ValueError(f"Strategy type '{strategy_type}' is not registered.")
        return strategy_class(aircraft_id=aircraft_id)
    
    def choice_policy(self,choice_type = "lasted") -> PPO:
        if choice_type == "lasted":
            self.current_policy = self.policy_manager.get_latest_policy()
        elif choice_type == "random":
            self.current_policy = self.policy_manager.get_random_policy()
        elif choice_type == "best":
            self.current_policy = self.policy_manager.get_best_policy()
        else:
            raise ValueError(f"Invalid choice_type '{choice_type}'. Must be 'lasted', 'random', or 'best'.")
        return self.current_policy
    
    @abstractmethod
    def GetState(self,origindata):
        raise NotImplementedError("GetState method must be implemented in the subclass.")
    def GetAction(self, action, originaldata = None, use_current_action=True):
        raise NotImplementedError("GetAction method must be implemented in the subclass.")
    def Reset(self):
        raise NotImplementedError("Reset method must be implemented in the subclass.")
    def GetTerminate(self,origindata):
        raise NotImplementedError("GetTerminate method must be implemented in the subclass.")
    
# 子策略管理器
# 管理所有子策略的创建和更新
class SubstrategiesManager():
    def __init__(self,planes_id_list=None): 
        self.sub_strategies: Dict[int, SubStrategy] = {}
    
        for plane_id in planes_id_list:
            self._init_sub_strategy(plane_id)
                
    def _init_sub_strategy(self, plane_id: int, default_strategy_type='search') -> None:
        self.update_sub_strategy(plane_id, default_strategy_type, plane_id=plane_id)

    def update_sub_strategy(self, strategy_type, plane_id=None,**kwargs) -> None:
        self.sub_strategies[plane_id] = SubStrategy.creat_strategy(strategy_type, plane_id=plane_id)

    def get_sub_strategy(self, plane_id):
        return self.sub_strategies[plane_id] if id in self.sub_strategies else None
    
    def is_sub_strategy_terminated(self, plane_id):
        result = self.sub_strategies[plane_id].GetTerminate(origindata=None)  
        return result

# # 对手子策略选择器
# class OpponentSubStrategySelector():
#     def __init__(self, planes_id_list=None):
#         self.opponent_sub_strategies = {}
#         for plane_id in planes_id_list:
#             self._init_opponent_sub_strategy(plane_id)

